# coding: utf-8

# -------------------------------------------------------------------------------
# Name:         postgres_client.py
# Description:
# Author:       XiangjunZhao
# EMAIL:        2419352654@qq.com
# Date:         2020/9/29 17:30
# -------------------------------------------------------------------------------
import logging

import psycopg2
from DBUtils.PooledDB import PooledDB

logger = logging.getLogger(__name__)


class PostgresClient(object):
    """
    Postgres客户端
    """

    def __init__(self, host=None, port=None, user=None, password=None, database=None):
        """
        初始化Postgres客户端
        Args:
            host: 主机地址
            port: 端口
            user: 用户名
            password: 密码
            database: 数据库
        """
        # 数据库地址
        self.host = host
        # 数据库端口
        self.port = port or 5432
        # 数据库用户名
        self.user = user
        # 数据库密码
        self.password = password
        # 数据库名称
        self.database = database
        self.db_pool = PooledDB(psycopg2, mincached=1, maxcached=5, host=self.host,
                                port=self.port, user=self.user, password=self.password,
                                database=self.database, connect_timeout=60)

    def get_connet(self):
        """
        获取数据库连接
        Returns:

        """
        try:
            # 获取数据库连接
            return self.db_pool.connection()
        except BaseException as e:
            logger.warning("获取数据库连接失败，异常：{}".format(e))
            self.get_connet()

    def close_connect(self, conn):
        """
        关闭数据库连接
        Args:
            conn: 数据库连接

        Returns:

        """
        if conn:
            conn.close()

    def close(self):
        """
        关闭数据库连接池
        Returns:

        """
        try:
            self.db_pool.close()
        except Exception as e:
            logger.error('数据库断开连接异常，原因：{}'.format(str(e)))

    def query_one(self, conn, sql):
        """
        查询一条数据
        Args:
            conn: 数据库连接
            sql: 查询sql

        Returns: dict

        """

        # 获取数据库游标
        with conn.cursor() as cursor:
            # 执行查询sql语句
            cursor.execute(sql)
            # 获取sql查询的列名
            desc = [item[0] for item in cursor.description]
            # 获取sql查询的结果
            data = cursor.fetchone()
            return dict(zip(desc, data)) if data else dict()

    def query_many(self, conn, sql):
        """
        查询所有数据
        Args:
            conn: 数据库连接
            sql: 查询sql

        Returns: list

        """

        # 获取数据库游标
        with conn.cursor() as cursor:
            # 执行查询sql语句
            cursor.execute(sql)
            # 获取sql查询的列名
            desc = [item[0] for item in cursor.description]
            # 获取sql查询的结果
            data = cursor.fetchall()
            return [dict(zip(desc, item)) for item in data] if data else list()

    def execute_one(self, conn=None, sql=None):
        """
        执行单条 插入、更新、删除sql
        Args:
            conn: 数据库连接
            sql: 插入、更新、删除sql

        Returns:

        """
        # 获取数据库游标
        with conn.cursor() as cursor:
            try:
                cursor.execute(sql)
            except psycopg2.Error as e:
                # 执行插入、更新、删除sql时发生异常，事务回滚
                conn.rollback()
                logger.error('数据库执行SQL异常，原因：{}'.format(str(e)))
            else:
                # 成功执行插入、更新、删除sql后，提交事务
                conn.commit()

    def execute_many(self, conn=None, sql=None, datas=None):
        """
        执行单条 插入、更新、删除sql
        Args:
            conn: 数据库连接
            sql: 插入、更新、删除sql
            datas: 插入、更新、删除sql的参数，参数格式为元组列表，示例：[(1,2,3,4),(1,2,3,4)]

        Returns:

        """
        # 获取数据库游标
        with conn.cursor() as cursor:
            try:
                cursor.executemany(sql, datas)
            except psycopg2.Error as e:
                # 执行插入、更新、删除sql时发生异常，事务回滚
                conn.rollback()
                logger.error('数据库执行SQL异常，原因：{}'.format(str(e)))
            else:
                # 成功执行插入、更新、删除sql后，提交事务
                conn.commit()
